#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/9/1
# @Author  : YunZhen
# @File    : download.py
# @Software: PyCharm
"""
Python安装包下载
"""

import os
import asyncio
import aiohttp
import aiofiles
from tqdm.asyncio import tqdm
from loguru import logger
from collections import defaultdict

from src.script._load import get_cached_versions
from src.script.verify import verify_download


async def get_download_links(version):
    """获取指定版本的所有下载链接"""

    # 从缓存获取版本下载连接
    all_python_versions = await get_cached_versions()
    if not all_python_versions:
        raise ValueError("无法获取版本缓存，确保已正确加载版本信息")

    download_links = defaultdict(list)
    for system, versions in all_python_versions.items():
        version_item = versions.get(version)
        if version_item:
            for name, link in version_item['links'].items():
                download_links[system].append(
                    {
                        'name': name,
                        'url': link,
                        'platform': system
                    }
                )
        else:
            logger.warning(f"版本 {version} 在系统 {system} 中未找到")
    logger.info(f"版本 {version} 的下载链接: {dict(download_links)}")
    return download_links


async def download_file(session, url, save_path):
    """下载文件（支持断点续传）"""
    # 检查已下载部分
    downloaded = 0
    if os.path.exists(save_path):
        downloaded = os.path.getsize(save_path)

    headers = {'Range': f'bytes={downloaded}-'} if downloaded else {}

    async with session.get(url, headers=headers) as response:
        # 验证是否支持断点续传
        if downloaded and response.status != 206:
            print("服务器不支持断点续传，重新开始下载")
            downloaded = 0
            headers = {}
            async with session.get(url) as new_response:
                response = new_response

        total_size = int(response.headers.get('content-length', 0)) + downloaded

        # 进度条设置
        file_name = os.path.basename(save_path)
        progress = tqdm(
            total=total_size,
            unit='B',
            unit_scale=True,
            desc=file_name,
            initial=downloaded
        )

        # 创建目录（如果不存在）
        os.makedirs(os.path.dirname(save_path), exist_ok=True)

        # 异步写入文件
        async with aiofiles.open(save_path, 'ab' if downloaded else 'wb') as f:
            while True:
                chunk = await response.content.read(1024 * 8)
                if not chunk:
                    break
                await f.write(chunk)
                progress.update(len(chunk))
        progress.close()

    # 校验文件完整性
    return await verify_download(save_path, total_size)


async def download_python_version(version: str, platforms: list, save_dir="packages"):
    """下载指定版本的Python安装包"""
    logger.info(f"获取 {version} 下载链接...")
    download_links = await get_download_links(version)

    if not download_links:
        logger.warning(f"未找到版本 {version} 的下载链接")
        return False

    async with aiohttp.ClientSession() as session:
        tasks = []
        for plat in platforms:
            if plat in download_links:
                for file_info in download_links[plat]:
                    filename = file_info["url"].split('/')[-1]
                    save_path = os.path.join(save_dir, plat, version, filename)
                    # 检查文件是否已完整存在
                    if os.path.exists(save_path):
                        file_size = os.path.getsize(save_path)
                        async with session.head(file_info["url"]) as response:
                            if response.status == 200:
                                total_size = int(response.headers.get('content-length', 0))
                                if file_size == total_size:
                                    print(f"文件已存在且完整: {file_info['name']}")
                                    continue

                    print(f"开始下载: {file_info['name']} -> {save_path}")
                    tasks.append(download_file(session, file_info["url"], save_path))
            else:
                print(f"版本 {version} 没有 {plat} 平台的安装包")

        # 并行下载
        if tasks:
            results = await asyncio.gather(*tasks)
            return all(results)
        else:
            print(f"所有文件已存在，跳过下载")
            return True
